#include "common.h"
#include "kernel_operator.h"
#include <cstdint>
namespace AscendC
{
template <typename _DT_X1_1, typename _DT_X1_2, typename _DT_X1_3, typename _DT_X2_1, typename _DT_X2_2,
          typename _DT_X2_3, typename _DT_Y_1, typename _DT_Y_2, typename _DT_Y_3>
class Kernel
{
  public:
    TPipe pipe;
    int64_t last_size = 0;
    CommonTiling tilingData;
    DefInTensor(X1_1);
    DefInTensor(X1_2);
    DefInTensor(X1_3);
    DefInTensor(X2_1);
    DefInTensor(X2_2);
    DefInTensor(X2_3);
    DefOutTensor(Y_1);
    DefOutTensor(Y_2);
    DefOutTensor(Y_3);

    DefBufVECIN(TEMP_1);
    DefBufVECIN(TEMP_2);

  public:
    __aicore__ inline Kernel()
    {
    }
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, CommonTiling &tiling)
    {
        tilingData = tiling;
        auto psize = tilingData.size / 3;
        last_size = psize % tilingData.bufferSize;
        if (last_size == 0)
        {
            last_size = tilingData.bufferSize;
        }
        // 切分global
        SetGlobalBuffer(X1_1, x1, tilingData.startIdx + 0 * psize, tilingData.blockLength);
        SetGlobalBuffer(X1_2, x1, tilingData.startIdx + 1 * psize, tilingData.blockLength);
        SetGlobalBuffer(X1_3, x1, tilingData.startIdx + 2 * psize, tilingData.blockLength);
        SetGlobalBuffer(X2_1, x2, tilingData.startIdx + 0 * psize, tilingData.blockLength);
        SetGlobalBuffer(X2_2, x2, tilingData.startIdx + 1 * psize, tilingData.blockLength);
        SetGlobalBuffer(X2_3, x2, tilingData.startIdx + 2 * psize, tilingData.blockLength);
        SetGlobalBuffer(Y_1, y, tilingData.startIdx + 0 * psize, tilingData.blockLength);
        SetGlobalBuffer(Y_2, y, tilingData.startIdx + 1 * psize, tilingData.blockLength);
        SetGlobalBuffer(Y_3, y, tilingData.startIdx + 2 * psize, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(X1_1, tilingData.bufferSize);
        InitQueueSimple(X1_2, tilingData.bufferSize);
        InitQueueSimple(X1_3, tilingData.bufferSize);
        InitQueueSimple(X2_1, tilingData.bufferSize);
        InitQueueSimple(X2_2, tilingData.bufferSize);
        InitQueueSimple(X2_3, tilingData.bufferSize);
        InitQueueSimple(Y_1, tilingData.bufferSize);
        InitQueueSimple(Y_2, tilingData.bufferSize);
        InitQueueSimple(Y_3, tilingData.bufferSize);

        InitTBufBuffer(TEMP_1, tilingData.bufferSize * sizeof(TypeOf(X1_1)));
        InitTBufBuffer(TEMP_2, tilingData.bufferSize * sizeof(TypeOf(X1_1)));
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn1(i, tilingData.bufferSize);
            Compute1(i, tilingData.bufferSize);
            CopyOut1(i, tilingData.bufferSize);
        }
        auto copy_count = ALIGN_TO(last_size, 32 / sizeof(TypeOf(X1_1)));
        CopyIn1(finnal_progress, copy_count);
        Compute1(finnal_progress, tilingData.finalCalcCount);
        CopyOut1(finnal_progress, copy_count);
    }
    __aicore__ inline void CopyIn1(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(X1_1, index, calcCount);
        EnQueGlobal2Local(X1_2, index, calcCount);
        EnQueGlobal2Local(X1_3, index, calcCount);
        EnQueGlobal2Local(X2_1, index, calcCount);
        EnQueGlobal2Local(X2_2, index, calcCount);
        EnQueGlobal2Local(X2_3, index, calcCount);
    }
    __aicore__ inline void Compute1(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(X1_1);
        DeQueSimple(X1_2);
        DeQueSimple(X1_3);
        DeQueSimple(X2_1);
        DeQueSimple(X2_2);
        DeQueSimple(X2_3);
        QueAllocSimple(Y_1);
        QueAllocSimple(Y_2);
        QueAllocSimple(Y_3);
        TBufGet(TEMP_1, TypeOf(X1_1));
        TBufGet(TEMP_2, TypeOf(X1_1));
        Mul(TBufTensorName(TEMP_1), LTensorName(X1_2), LTensorName(X2_3), calcCount);
        Mul(TBufTensorName(TEMP_2), LTensorName(X1_3), LTensorName(X2_2), calcCount);
        Sub(LTensorName(Y_1), TBufTensorName(TEMP_1), TBufTensorName(TEMP_2), calcCount);

        Mul(TBufTensorName(TEMP_1), LTensorName(X1_3), LTensorName(X2_1), calcCount);
        Mul(TBufTensorName(TEMP_2), LTensorName(X1_1), LTensorName(X2_3), calcCount);
        Sub(LTensorName(Y_2), TBufTensorName(TEMP_1), TBufTensorName(TEMP_2), calcCount);

        Mul(TBufTensorName(TEMP_1), LTensorName(X1_1), LTensorName(X2_2), calcCount);
        Mul(TBufTensorName(TEMP_2), LTensorName(X1_2), LTensorName(X2_1), calcCount);
        Sub(LTensorName(Y_3), TBufTensorName(TEMP_1), TBufTensorName(TEMP_2), calcCount);
        EnQue(Y_1);
        EnQue(Y_2);
        EnQue(Y_3);
        QueFree(X1_1);
        QueFree(X1_2);
        QueFree(X1_3);
        QueFree(X2_1);
        QueFree(X2_2);
        QueFree(X2_3);
    }
    __aicore__ inline void CopyOut1(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y_1, index, calcCount);
        DeQueLocal2Global(Y_2, index, calcCount);
        DeQueLocal2Global(Y_3, index, calcCount);
    }
};
} // namespace AscendC

extern "C" __aicore__ inline void trans_data(GM_ADDR x, GM_ADDR y, int64_t presize, int64_t stride, int64_t size)
{
    auto x1_casted = CAST_GM_ADDR(DTYPE_X1, x, 0);
    auto y_casted = CAST_GM_ADDR(DTYPE_Y, y, 0);
    for (auto i = 0; i < presize; ++i)
    {
        for (auto j = 0; j < 3; ++j)
        {
            for (auto k = 0; k < stride; ++k)
            {
                y_casted[j * presize * stride + i * stride + k] = x1_casted[i * 3 * stride + j * stride + k];
            }
        }
    }

    for (auto i = 0; i < size; ++i)
    {
        x1_casted[i] = y_casted[i];
    }
}

extern "C" __global__ __aicore__ void cross(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    // 处理输入
    auto presize = tiling_data.preSize;
    auto stride = tiling_data.stride;
    auto p_size = presize * stride;
    trans_data(x1, y, presize, stride, tiling_data.size);
    trans_data(x2, y, presize, stride, tiling_data.size);

    switch (tiling_data.dType)
    {
    // float
    case 0: {
        AscendC::Kernel<float, float, float, float, float, float, float, float, float> op;
        op.Init(x1, x2, y,  commonTiling);
        op.Process();
    }
    break;
    // half
    case 1: {
        AscendC::Kernel<half, half, half, half, half, half, half, half, half> op;
        op.Init(x1, x2, y,  commonTiling);
        op.Process();
    }
    break;
    // int8_t
    case 2: {
        // AscendC::KernelInt8<int8_t, int8_t, DTYPE_Y> op;
        // op.Init(x1, x2, y, commonTiling);
        // op.Process();
    }
    break;
    // int32_t
    case 3: {
        AscendC::Kernel<int32_t, int32_t, int32_t, int32_t, int32_t, int32_t, int32_t, int32_t, int32_t> op;
        op.Init(x1, x2, y,  commonTiling);
        op.Process();
    }
    break;
    }

    // 处理结果
    auto x1_casted = CAST_GM_ADDR(DTYPE_X1, x1, 0);
    auto y_casted = CAST_GM_ADDR(DTYPE_Y, y, 0);
    for (auto i = 0; i < presize; ++i)
    {
        for (auto j = 0; j < 3; ++j)
        {
            for (auto k = 0; k < stride; ++k)
            {
                x1_casted[i * 3 * stride + j * stride + k] = y_casted[j * presize * stride + i * stride + k];
            }
        }
    }

    for (auto i = 0; i < tiling_data.size; ++i)
    {
        y_casted[i] = x1_casted[i];
    }
}